import tensorflow as tf
import numpy as np
from cn.redguest.pbase.model.Dense import Dense as Dense
from cn.redguest.pbase.model.Processor import *


def main():
    bat_size = 1000
    train_num = 30000
    learning_speed = 1e-2
    tf_number_type = tf.float64
    np_number_type = np.float64

    x = tf.placeholder(tf_number_type)
    y = tf.placeholder(tf_number_type)

    l1 = Dense(x, 1, 2, tf.nn.tanh, d_type=tf_number_type)
    l2 = Dense(l1.y, 2, 3, tf.nn.tanh, d_type=tf_number_type)
    l3 = Dense(l2.y, 3, 4, tf.nn.tanh, d_type=tf_number_type)
    l4 = Dense(l3.y, 4, 5, tf.nn.relu, d_type=tf_number_type)
    l5 = Dense(l4.y, 5, 5, d_type=tf_number_type)

    loss = tf.reduce_mean(tf.square(y - l5.y))

    r1 = L2Regularizer(l1.Weights, 0.01)
    r2 = L2Regularizer(l2.Weights, 0.01)
    r3 = L2Regularizer(l3.Weights, 0.01)
    r4 = L2Regularizer(l4.Weights, 0.01)
    r5 = L2Regularizer(l5.Weights, 1e-2)

    loss += r1.y + r2.y + r3.y + r4.y + r5.y

    optimizer = tf.train.GradientDescentOptimizer(learning_speed)
    train = optimizer.minimize(loss)

    init = tf.initialize_all_variables()

    s = tf.Session()
    s.run(init)

    merged = tf.summary.merge_all()
    writer = tf.summary.FileWriter('logs', s.graph, flush_secs=10)

    train_x = [[0.1], [0.3], [0.4]]
    train_y = [[0.1, 0.5, 0.4, 0.2, 0.1], [0.1, 0.5, 0.8, 0.6, 0.1], [0.1, 0.5, 0.9, 0.9, 0.1]]

    for i in range(train_num):
        _ = s.run(train, feed_dict={
            x: train_x,
            y: train_y
        })
        if i % 1000 == 0:
            print(l1.Weights.eval(s), l2.Weights.eval(s), l3.Weights.eval(s))

    print(np.round(np.multiply(s.run(l5.y, feed_dict={x: [[0.2]]}), 10)))


if __name__ == "__main__":
    main()
